package com.pg.agent.engine.service.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.pg.agent.engine.client.message.InPutMessage;
import com.pg.agent.engine.client.model.SessionObject;
import com.pg.agent.engine.client.model.PromptConstant;
import com.pg.agent.engine.client.model.agent.AIAgent;
import com.pg.agent.engine.client.model.environment.AIEnvironment;
import com.pg.agent.engine.client.model.llm.CallLog;
import com.pg.agent.engine.client.model.llm.ComponentLLM;
import com.pg.agent.engine.client.model.sop.AISOP;
import com.pg.agent.engine.client.model.sop.Memory;
import com.pg.agent.engine.service.service.AIAgentService;
import com.pg.agent.engine.service.service.AIEnvironmentService;
import com.pg.agent.engine.service.utils.ComponentLLMUtil;
import com.pg.agent.engine.service.utils.MilvusUtils;
import com.pg.base.util.PgHttpUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;


@Service
@Slf4j
public class AIEnvironmentServiceImpl implements AIEnvironmentService {
    private static final int EXPIRE_TIME = 24 * 60 * 60 * 1000;
    @Resource
    private MilvusUtils milvusUtils;

/*    @Resource
    private AIAgentService aiAgentService;*/

    @Resource
    private Executor taskExecutor;

    @Value("${embedding.url:null}")
    private String embeddingUrl;

//    @Override
//    public void updateMemory(AIAgentSession session, Memory memory, ManagerObject managerObject, Long saasAppId, Integer exeType, Integer type) {
//
//        AIEnvironment aiEnvironment = session.getEnvironment();
//        Integer maxChatHistory = 10;
//        // 把记忆添加到environment 长记忆中
//        List<Memory> longTermMemory = aiEnvironment.getShareMemory().getLongMemory();
//        longTermMemory.add(memory);
//        JedisUtils.setObject(session.getUserKey(), managerObject, EXPIRE_TIME);
//
//        // text-embedding-ada-002  gpt-35-turbo-16k
//        ComponentLLM llm = session.getAgent().getComLLM();
//        taskExecutor.execute(() -> {
//            //计算向量值
//            List<Double> currentEmbedding = getEmbedding(session.getEnvironment(), memory.getContent());
//            if (currentEmbedding != null) {
//                memory.setChatEmbeddings(currentEmbedding);
//
//                milvusUtils.insertData(managerObject.getUserId(), managerObject.getSessionId(),
//                        session.getSop().getId(), session.getState().getName(), memory.getSendName(), memory.getSendRole() == null ? "": memory.getSendName(), memory.getContent(), currentEmbedding, saasAppId, exeType, type);
//
//                // 更新对应agent记忆
//                aiAgentService.updateMemory(session.getAgent(), session, memory);
//
//                //长记忆取余。统计后，赋值给短记忆。
//                if (longTermMemory.size() % maxChatHistory == 0) {
//                    //String shortSummary = summary(session.getEnvironment(), session, llm);
//                    //aiEnvironment.getShareMemory().setShortMemory(shortSummary);
//                }
//
//                JedisUtils.setObject(session.getUserKey(), managerObject, EXPIRE_TIME);
//            }
//        });
//
//
//    }

    @Override
    public void updateMemory(InPutMessage inPutMessage, Memory memory, SessionObject sessionObject) {
        AIEnvironment aiEnvironment = inPutMessage.getEnvironment();
        Integer maxChatHistory = 10;
        // 把记忆添加到environment 长记忆中
        List<Memory> longTermMemory = aiEnvironment.getShareMemory().getLongMemory();
        longTermMemory.add(memory);
        //JedisUtils.set(inPutMessage.getUserKey(),  JSONObject.toJSONString(managerObject), EXPIRE_TIME);

        // text-embedding-ada-002  gpt-35-turbo-16k
        ComponentLLM llm = inPutMessage.getSop().getComLLM();
        taskExecutor.execute(() -> {
            //计算向量值
            List<Double> currentEmbedding = getEmbedding(inPutMessage.getEnvironment(), memory.getContent());
            if (currentEmbedding != null) {
                memory.setChatEmbeddings(currentEmbedding);
//
//                milvusUtils.insertData(managerObject.getUserId(), managerObject.getSessionId(),
//                        inPutMessage.getSop().getId(), inPutMessage.getState().getName(),
//                        memory.getSendName(), memory.getSendRole() == null ? "": memory.getSendName(), memory.getContent(),
//                        currentEmbedding, saasAppId, exeType, type);
               //milvusUtils.insertMemory(managerObject, inPutMessage, memory);

                // 更新对应agent记忆
                if (inPutMessage.getAgent() != null) {
                    //aiAgentService.updateMemory(inPutMessage, memory);
                }

                //长记忆取余。统计后，赋值给短记忆。
                if (longTermMemory.size() % maxChatHistory == 0) {
                    AISOP aisop = inPutMessage.getSop();
                    if (aisop.isSummaryHistory()) {
                        String shortSummary = summary(inPutMessage.getEnvironment(), inPutMessage, llm);
                        aiEnvironment.getShareMemory().setShortMemory(shortSummary);
                    }
                }

                //JedisUtils.set(inPutMessage.getUserKey(),  JSONObject.toJSONString(managerObject), EXPIRE_TIME);
            }
        });
    }

    @Override
    public List<Double> getEmbedding(AIEnvironment environment, String content) {
/*        JSONObject jsonObject = new JSONObject();
        jsonObject.put("content", content);
        String resultStr = PgHttpUtil.postBody(embeddingUrl, null,jsonObject.toJSONString());
        JSONObject result = JSONObject.parseObject(resultStr);
        if (result != null && ObjectUtil.equal(result.getInteger("code"), 0)) {
            JSONObject data = (JSONObject) result.getJSONObject("data").getJSONArray("data").get(0);
            List<Double> list = new ArrayList<>(); // 创建空列表
            JSONArray array = data.getJSONArray("embedding");
            for (int i = 0; i < array.size(); i++) {
                double value = array.getDoubleValue(i); // 逐个提取元素并转换为 Double 类型
                list.add(value); // 添加到列表中
            }
            return list;
        }*/
        return null;


    }

    @Override
    public String summary(AIEnvironment aiEnvironment, InPutMessage message, ComponentLLM llm) {

        //
        List<Memory> longTermMemory = aiEnvironment.getShareMemory().getLongMemory();

        String relevantMemory = "";
        Memory memory = longTermMemory.get(longTermMemory.size() - 1);
        List<Double> chatEmbedding = memory.getChatEmbeddings();
        String expr = "session_id == '" + message.getSessionId() + "'";
        if (chatEmbedding != null && chatEmbedding.size() > 0) {
            long search = System.currentTimeMillis();
            //relevantMemory = milvusUtils.getRelevantHistory(chatEmbedding, expr, 3, null, 10);
            log.info("调用查询用时：{}", System.currentTimeMillis() - search);
        } else {
            log.info("===============需要获取embedding===============");
            long start = System.currentTimeMillis();
            List<Double> currentEmbedding = getEmbedding(aiEnvironment, memory.getContent());
            log.info("调用embedding用时：{}", System.currentTimeMillis() - start);
            if (currentEmbedding != null) {
                long search = System.currentTimeMillis();
                //relevantMemory = milvusUtils.getRelevantHistory(currentEmbedding, expr, 3, null, 10);
                memory.setChatEmbeddings(currentEmbedding);
                log.info("调用查询用时：{}", System.currentTimeMillis() - search);
            }
        }
        // 组装长期记忆为字符串。
        String chatHistory = Memory.getChatHistory(longTermMemory);
        String summary = aiEnvironment.getShareMemory().getShortMemory();

        String environmentPrompt = StrUtil.trimToEmpty(aiEnvironment.getEnvironmentPrompt().get(message.getState().getName()));
        String summarySystemPrompt = StrUtil.trimToEmpty(aiEnvironment.getSummarySystemPrompt().get(message.getState().getName()));
        String currentMemory = StrUtil.trimToEmpty(String.format(PromptConstant.ENVIRONMENT_SUMMARY_MEMORY, summary, chatHistory, relevantMemory));

        //系统总结提示语
        String systemSummaryPrompt = String.format(PromptConstant.ENVIRONMENT_SUMMARY_SYSTEM_PROMPT, environmentPrompt, currentMemory, summarySystemPrompt);

        //调用大模型
        CallLog callLog = new CallLog(message.getState().getName(), "上下文总结");


        ComponentLLMUtil.call(callLog, llm, message, null, systemSummaryPrompt, "");

        return callLog.getResponse();
    }

    @Override
    public Integer getAgentLastConversationIdx(AIAgent agent, List<Memory> currentLongTermMemoryList) {

        Integer lastConversationIdx = -1;

        if (CollectionUtil.isNotEmpty(currentLongTermMemoryList)) {
            for (int i = 0; i < currentLongTermMemoryList.size(); i++) {
                if (ObjectUtil.isAllNotEmpty(currentLongTermMemoryList.get(i), agent)) {
                    if (ObjectUtil.equal(currentLongTermMemoryList.get(i).getSendName(), agent.getName())) {
                        lastConversationIdx = i;
                    }
                }
            }
        }
        return lastConversationIdx;
    }

    @Override
    public Memory observe(AIEnvironment environment, AIAgent agent, InPutMessage message) {

        // 发送的最大聊天次数
        Integer maxChatHistory = environment.getConfig().getMaxChatHistory();

        //判断环境类型EnvironmentType
        //cooperative:在不同的状态之间共享信息;
        //competive :不同状态之间没有信息共享
        Integer currentChatHistoryIdx = "competive".equals(environment.getEnvironmentType()) ? environment.getCurrentChatHistoryIdx() : 0;

        //上下文长期记忆和向量数据
        List<Memory> envLongMemory = environment.getShareMemory().getLongMemory().subList(currentChatHistoryIdx, environment.getShareMemory().getLongMemory().size());
        List<List<Double>> embeddings = envLongMemory.stream().map(e -> e.getChatEmbeddings()).collect(Collectors.toList());

        // 对长期记忆数据量处理，判断是否大于最大聊天数，如果大于则截取最后的最大聊天数量集合
        if (ObjectUtil.isNotEmpty(envLongMemory) && envLongMemory.size() > maxChatHistory) {
            envLongMemory = envLongMemory.subList(envLongMemory.size() - maxChatHistory, envLongMemory.size());
            embeddings = embeddings.subList(0, maxChatHistory + 1);
        }

        //agent 获取相关记忆数据。
        String relevantMemory = "";

        if (ObjectUtil.isNotEmpty(envLongMemory) && envLongMemory.size() > 10) {

//            Memory memory = envLongMemory.get(envLongMemory.size() - 1);
//            String query = (memory == null) ? "" : memory.getContent();
//            List<Memory> relevantHistory = Utils.getRelevantHistory(query, envLongMemory, embeddings, environment.getConfig());
//            relevantMemory = Memory.getChatHistory(relevantHistory, agent.getName());
            Memory memory = envLongMemory.get(envLongMemory.size() - 1);
            List<Double> chatEmbedding = memory.getChatEmbeddings();
            String expr = "session_id == '" + message.getSessionId() + "'";
            if (chatEmbedding != null && chatEmbedding.size() > 0) {
                long search = System.currentTimeMillis();
                //relevantMemory = milvusUtils.getRelevantHistory(chatEmbedding, expr, 3, null, 10);
                log.info("调用查询用时：{}", System.currentTimeMillis() - search);
            } else {
                log.info("===============需要获取embedding===============");
                long start = System.currentTimeMillis();
                List<Double> currentEmbedding = getEmbedding(environment, memory.getContent());
                log.info("调用embedding用时：{}", System.currentTimeMillis() - start);
                if (currentEmbedding != null) {
                    long search = System.currentTimeMillis();
                    //relevantMemory = milvusUtils.getRelevantHistory(currentEmbedding, expr, 3, null, 10);
                    memory.setChatEmbeddings(currentEmbedding);
                    log.info("调用查询用时：{}", System.currentTimeMillis() - search);
                }
            }

        }

        //relevantMemory = String.format(PromptConstant.AGENT_OBSERVE_RELEVANT_MEMORY, chatHistoryStr.toString());
        // 存储到agent RelevantMemory上。
        agent.setRelevantMemory(relevantMemory);

        // 组装历史对象记录
        String agentChatHistory = Memory.getChatHistory(envLongMemory);

        //获取最新（最后）一条对话
        String newConversations = getAgentNewMemory(envLongMemory);

        //组装记忆内容
        String currentMemoryContent = String.format(PromptConstant.AGENT_OBSERVE_MEMORY, StrUtil.trimToEmpty(agent.getShortTermMemory()), relevantMemory, agentChatHistory, newConversations);

        // 返回记忆对象
        return new Memory(message.getExeId(), message.getDataId(), "user", "", currentMemoryContent);
    }

    @Override
    public String getAgentNewMemory(List<Memory> currentLongTermMemoryList) {

        //Integer lastConversationIdx = getAgentLastConversationIdx(agent, currentLongTermMemoryList);

        List<Memory> newConversation = new ArrayList<>();

        if (ObjectUtil.isNotEmpty(currentLongTermMemoryList)) {
            newConversation = currentLongTermMemoryList.subList(currentLongTermMemoryList.size() - 1, currentLongTermMemoryList.size());
        }

        return Memory.getChatHistory(newConversation);
    }

}
